#include <stdio.h>
#include <cuda.h>
#include <curand_kernel.h>
#include <cuda_runtime.h>
#include "GPUPoissonCuda.h"
// For the CUDA runtime routines (prefixed with "cuda_")
#include "cuda_errorcheck.h"
#include "cuda_defs.h"
#include "cuda_reduceMin.h"
#include "gPlannerCollisionWrapper.h"
#include "cuda_statePrimitives.h"
#include "cuda_stateIntersect.h"
#include "cuda_rrt_types.h"
#include "GPUMemoryManager.h"

// pointers for data on device memory
////////////////////////////////////////////////////////////////////////////////
extern curandState* d_randStates; // M * 3

GMesh* d_robot;
GMesh* d_obstacle;

__device__ unsigned int d_numSamples;
PoissonDiskSampleData d_sampleData;

RrtTree d_startTree;
RrtTree d_goalTree;

//SerializedStates* d_startState;
//SerializedStates* d_goalState;

// data on host memory
////////////////////////////////////////////////////////////////////////////////
unsigned int h_numPrecomputedSamples = 0;
unsigned int h_numSamples = 0;
unsigned int h_numStartTreeNodes = 0;
unsigned int h_numGoalTreeNodes = 0;

// device memory constants
////////////////////////////////////////////////////////////////////////////////
__device__ __constant__ float c_boundsMin[3];
__device__ __constant__ float c_boundsMax[3];
__device__ __constant__ float c_boundsDiff[3];
__device__ __constant__ float c_radius[MAX_LEVEL];



// poisson-disk sample data
float* d_samples; // 3(translation) + 4(rotation) floats
int* d_levels;
int* d_isTreeNode;

// start tree
unsigned int* d_numStartTreeNodes;
int* d_startTreeNodeIndices; // index to samples
int* d_startTreeNodeParentIndices; // index to d_startTreeNodeXXX
float* d_startTreeNodes;

// goal tree
unsigned int* d_numGoalTreeNodes;
int* d_goalTreeNodeIndices;
int* d_goalTreeNodeParentIndices;
float* d_goalTreeNodes;

// start/goal states
float* d_startPos;
float* d_goalPos;

// internal states
float* d_internalStates;	// M
float* d_internalStates2;	// M

// internal distances
float* d_distances;			// M * (K + 512)
float* d_internalDistances; // M * MAX_SAMPLES

// internal indices
int* d_nnIndices;			// M * K
int* d_nnIndices2;			// M * K
int* d_nnIndices3;			// M * (K + 512)

// internal value
int* d_internalValues;		// M
int* d_internalValues2;		// M

int* d_result;

__global__ void reduceMin1(float *g_idata, float *g_odata, unsigned int query_width, int* g_oIndices, bool removeMin)
{
	extern __shared__ float sdata[];

	unsigned int queryIdx = blockIdx.y;
	unsigned int dataIdx = queryIdx;

	g_odata[queryIdx] = g_idata[queryIdx];
	g_oIndices[queryIdx] = 0;

	printf("min : %f at %d is written to %x\n", g_idata[queryIdx], 0, g_oIndices + queryIdx);

	if (removeMin)
		g_idata[dataIdx] = 1E+37;
}

__global__ void reduceMin1WithIndices(float *g_idata, int* g_iIndices, float *g_odata, unsigned int n, unsigned int query_width, int* g_oIndices, bool removeMin)
{
	extern __shared__ float sdata[];

	unsigned int queryIdx = blockIdx.y;
	unsigned int refIdx = g_iIndices[queryIdx];
	unsigned int dataIdx = queryIdx * n + refIdx;

	g_odata[queryIdx] = g_idata[dataIdx];
	g_oIndices[queryIdx] = dataIdx;

	if (removeMin)
		g_idata[dataIdx] = 1E+37;
}

__global__ void computeDistanceMatrix(float* ref, int ref_width, float* query, int query_width, float* dist)
{
	__shared__ float queryState[STATE_STRIDE];
	__shared__ float squared[STATE_STRIDE * 32];

	int queryIdx = blockIdx.y;
	int refIdx = blockIdx.x * 32 + threadIdx.y;
	int componentIdx = threadIdx.x;

	float q, r, x1, x2;

	if (refIdx < ref_width)
	{
		// read from global memory
		if (threadIdx.y == 0)
		{
			queryState[componentIdx] = query[STATE_STRIDE * queryIdx + componentIdx];
		}
		__syncthreads();

		q = queryState[componentIdx];
		r = ref[STATE_STRIDE * refIdx + componentIdx];

		if (componentIdx < 3)
		{
			squared[STATE_STRIDE * threadIdx.y + componentIdx] = (q - r) * (q - r);
		}
		else
		{
			squared[STATE_STRIDE * threadIdx.y + componentIdx] = q * r;
		}
		__syncthreads();

		if (componentIdx == 0)
		{
			x1 = sqrtf(squared[STATE_STRIDE * threadIdx.y + 0] + squared[STATE_STRIDE * threadIdx.y + 1] + squared[STATE_STRIDE * threadIdx.y + 2]);
			x2 = fabs(squared[STATE_STRIDE * threadIdx.y + 3] + squared[STATE_STRIDE * threadIdx.y + 4] +
				squared[STATE_STRIDE * threadIdx.y + 5] + squared[STATE_STRIDE * threadIdx.y + 6]);
			x2 = (x2 > 1 - 1e-7) ? 0.0 : acosf(x2);
			dist[queryIdx * ref_width + refIdx] = x1 + x2;

			{
				printf("Dist btn (%.2f %.2f %.2f)%d (%.2f %.2f %.2f)%d : %.2f %.2f\n", 
					query[STATE_STRIDE * queryIdx + 0],
					query[STATE_STRIDE * queryIdx + 1],
					query[STATE_STRIDE * queryIdx + 2],
					queryIdx, 
					ref[STATE_STRIDE * refIdx + 0],
					ref[STATE_STRIDE * refIdx + 1],
					ref[STATE_STRIDE * refIdx + 2],
					refIdx, x1, x2);
			}
		}
	}
}

void knn(float* ref, int ref_width, float* query, int query_width, unsigned int k, float* dist, int* ind)
{
	int gridX = (ref_width + 31) / 32;
	dim3 grids = dim3(gridX, query_width, 1);
	dim3 threads = dim3(7, 32, 1);

	computeDistanceMatrix <<< grids, threads >>> (ref, ref_width, query, query_width, d_internalDistances);

	float* distWrite = dist;
	int* indWrite = ind;
	for (int i = 0; i < k; ++i)
	{
		if (ref_width > 512 * 512)
		{
			// TODO: 3-level
		}
		else if (ref_width > 512)
		{
			int blockSize = 512;
			gridX = (ref_width + (blockSize-1)) / blockSize;
			grids = dim3(gridX, query_width, 1);
			threads = dim3(blockSize, 1, 1);
			int smemSize = blockSize * (sizeof(float) + sizeof(int));

			float* distWrite2 = distWrite + query_width;
			int* indWrite2 = indWrite + query_width;
			reduceMin<256, float><<< grids, threads, smemSize >>>(d_internalDistances, distWrite2, ref_width, query_width, indWrite2, false);

			// distWrite2, indWrite2 has gridX * query_width results;

			int nextStepWidth = gridX;
			while (blockSize / 2 >= nextStepWidth)
			{
				blockSize /= 2;
			}
			gridX = (nextStepWidth + (blockSize-1)) / blockSize;
			grids = dim3(gridX, query_width, 1);
			threads = dim3(blockSize, 1, 1);
			smemSize = blockSize * (sizeof(float) + sizeof(int));

			switch (blockSize)
			{
			case 512:
				reduceMinWithIndices<256, float><<< grids, threads, smemSize >>>(d_internalDistances, indWrite2, distWrite, ref_width, nextStepWidth, query_width, indWrite, true); break;
			case 256:
				reduceMinWithIndices<128, float><<< grids, threads, smemSize >>>(d_internalDistances, indWrite2, distWrite, ref_width, nextStepWidth, query_width, indWrite, true); break;
			case 128:
				reduceMinWithIndices< 64, float><<< grids, threads, smemSize >>>(d_internalDistances, indWrite2, distWrite, ref_width, nextStepWidth, query_width, indWrite, true); break;
			case 64:
				reduceMinWithIndices< 32, float><<< grids, threads, smemSize >>>(d_internalDistances, indWrite2, distWrite, ref_width, nextStepWidth, query_width, indWrite, true); break;
			case 32:
				reduceMinWithIndices< 16, float><<< grids, threads, smemSize >>>(d_internalDistances, indWrite2, distWrite, ref_width, nextStepWidth, query_width, indWrite, true); break;
			case 16:
				reduceMinWithIndices< 8, float><<< grids, threads, smemSize >>>(d_internalDistances, indWrite2, distWrite, ref_width, nextStepWidth, query_width, indWrite, true); break;
			case 8:
				reduceMinWithIndices< 4, float><<< grids, threads, smemSize >>>(d_internalDistances, indWrite2, distWrite, ref_width, nextStepWidth, query_width, indWrite, true); break;
			case 4:
				reduceMinWithIndices< 2, float><<< grids, threads, smemSize >>>(d_internalDistances, indWrite2, distWrite, ref_width, nextStepWidth, query_width, indWrite, true); break;
			case 2:
				reduceMinWithIndices< 1, float><<< grids, threads, smemSize >>>(d_internalDistances, indWrite2, distWrite, ref_width, nextStepWidth, query_width, indWrite, true); break;
			case 1:
				reduceMin1WithIndices<<< grids, threads, smemSize >>>(d_internalDistances, indWrite2, distWrite, ref_width, query_width, indWrite, true); break;
			}



		}
		else
		{
			int blockSize = 512;
			while (blockSize / 2 >= ref_width)
			{
				blockSize /= 2;
			}
			gridX = (ref_width + (blockSize-1)) / blockSize;
			grids = dim3(gridX, query_width, 1);
			threads = dim3(blockSize, 1, 1);
			int smemSize = blockSize * (sizeof(float) + sizeof(int));

			switch (blockSize)
			{
			case 512:
				reduceMin<256, float><<< grids, threads, smemSize >>>(d_internalDistances, distWrite, ref_width, query_width, indWrite, true); break;
			case 256:
				reduceMin<128, float><<< grids, threads, smemSize >>>(d_internalDistances, distWrite, ref_width, query_width, indWrite, true); break;
			case 128:
				reduceMin< 64, float><<< grids, threads, smemSize >>>(d_internalDistances, distWrite, ref_width, query_width, indWrite, true); break;
			case 64:
				reduceMin< 32, float><<< grids, threads, smemSize >>>(d_internalDistances, distWrite, ref_width, query_width, indWrite, true); break;
			case 32:
				reduceMin< 16, float><<< grids, threads, smemSize >>>(d_internalDistances, distWrite, ref_width, query_width, indWrite, true); break;
			case 16:
				reduceMin< 8, float><<< grids, threads, smemSize >>>(d_internalDistances, distWrite, ref_width, query_width, indWrite, true); break;
			case 8:
				reduceMin< 4, float><<< grids, threads, smemSize >>>(d_internalDistances, distWrite, ref_width, query_width, indWrite, true); break;
			case 4:
				reduceMin< 2, float><<< grids, threads, smemSize >>>(d_internalDistances, distWrite, ref_width, query_width, indWrite, true); break;
			case 2:
				reduceMin< 1, float><<< grids, threads, smemSize >>>(d_internalDistances, distWrite, ref_width, query_width, indWrite, true); break;
			case 1:
				reduceMin1<<< grids, threads, smemSize >>>(d_internalDistances, distWrite, query_width, indWrite, true); break;
			}
		}
		distWrite += query_width;
		indWrite += query_width;
	}
}

__global__ void computeDistancesTreeNodesToStatesKernel(int queryWidth, float* treeNodes, int* nodeIndices, float* destStates, 
	float* distR3, float* distSO3)
{
	__shared__ int nodeIndex[32];
	__shared__ float srcState[STATE_STRIDE * 32];
	__shared__ float destState[STATE_STRIDE * 32];
	__shared__ float squared[STATE_STRIDE * 32];

	int queryIdx = blockIdx.x * 32 + threadIdx.y;
	int componentIdx = threadIdx.x;
	int sharedDataIdx = threadIdx.y * STATE_STRIDE + threadIdx.x;

	float src, dest, x1, x2;

	if (queryIdx < queryWidth)
	{
		// read from global memory
		if (componentIdx == 0)
			nodeIndex[threadIdx.y] = nodeIndices[queryIdx];
		__syncthreads();

		int localIdx = nodeIndex[threadIdx.y];
		src = treeNodes[STATE_STRIDE * localIdx + componentIdx];
		srcState[sharedDataIdx] = src;
		dest = destStates[STATE_STRIDE * queryIdx + componentIdx];
		destState[sharedDataIdx] = dest;
		__syncthreads();

		if (componentIdx < 3)
		{
			squared[sharedDataIdx] = (dest - src) * (dest - src);
		}
		else
		{
			squared[sharedDataIdx] = src * dest;
		}
		__syncthreads();

		if (componentIdx == 0)
		{
			distR3[queryIdx] = sqrtf(squared[STATE_STRIDE * threadIdx.y + 0] + squared[STATE_STRIDE * threadIdx.y + 1] + squared[STATE_STRIDE * threadIdx.y + 2]);
			x2 = fabs(squared[STATE_STRIDE * threadIdx.y + 3] + squared[STATE_STRIDE * threadIdx.y + 4] +
				squared[STATE_STRIDE * threadIdx.y + 5] + squared[STATE_STRIDE * threadIdx.y + 6]);
			x2 = (x2 > 1 - 1e-7) ? 0.0 : acosf(x2);
			distSO3[queryIdx] = x2;

			printf("%d dist from (%.2f %.2f %.2f) to (%.2f %.2f %.2f) : (%f)\n", queryIdx,
				srcState[STATE_STRIDE * threadIdx.y + 0],
				srcState[STATE_STRIDE * threadIdx.y + 1],
				srcState[STATE_STRIDE * threadIdx.y + 2],
				destState[STATE_STRIDE * threadIdx.y + 0],
				destState[STATE_STRIDE * threadIdx.y + 1],
				destState[STATE_STRIDE * threadIdx.y + 2],
				distR3[queryIdx] + x2);
		}
	}
}

void computeDistancesTreeNodesToStates(int queryWidth, float* treeNodes, int* nodeIndices, float* destStates, float* distR3, float* distSO3)
{
	int gridX = (queryWidth + 31) / 32;
	dim3 grids = dim3(gridX, 1, 1);
	dim3 threads = dim3(7, 32, 1);

	computeDistancesTreeNodesToStatesKernel <<< grids, threads >>> (queryWidth, treeNodes, nodeIndices, destStates, distR3, distSO3);
}

__global__ void computeIntersectionStatesKernel(int queryWidth, float* treeNodes, int* nodeIndices, 
	float* destStates, float* distR3, float* distSO3, float* intersectionStates, 
	int* treeNodeSampleIndices, int* levels)
{
	__shared__ int nodeIndex[32];
	__shared__ float srcState[STATE_STRIDE * 32];
	__shared__ float destState[STATE_STRIDE * 32];
	__shared__ float distanceR3[32];
	__shared__ float distanceSO3[32];
	__shared__ float ratio[32];
	__shared__ float so3Inter[STATE_STRIDE * 32];
	__shared__ float so3s0[32];
	__shared__ float so3s1[32];
	__shared__ float so3d[32];

	int queryIdx = blockIdx.x * 32 + threadIdx.y;
	int componentIdx = threadIdx.x;
	int sharedDataIdx = threadIdx.y * STATE_STRIDE + threadIdx.x;

	float src, dest, x1, x2;

	if (queryIdx < queryWidth)
	{
		// read src/dest/distance from global memory
		if (componentIdx == 0)
		{
			nodeIndex[threadIdx.y] = nodeIndices[queryIdx];
			distanceR3[threadIdx.y] = distR3[queryIdx];
			distanceSO3[threadIdx.y] = distSO3[queryIdx];

			int level = levels[treeNodeSampleIndices[nodeIndex[threadIdx.y]]];
			ratio[threadIdx.y] = c_radius[level] / (distanceR3[threadIdx.y] + distanceSO3[threadIdx.y]);
		}

		dest = destStates[STATE_STRIDE * queryIdx + componentIdx];
		destState[sharedDataIdx] = dest;
		__syncthreads();

		int localIdx = nodeIndex[threadIdx.y];
		src = treeNodes[STATE_STRIDE * localIdx + componentIdx];
		srcState[sharedDataIdx] = src;
		__syncthreads();

		// compute interpolation pt
		if (componentIdx < 3)
		{
			intersectionStates[STATE_STRIDE * queryIdx + componentIdx] = src + (dest - src) * ratio[threadIdx.y];
		}
		else
		{
			// SO3 interpolation
			intersectionStates[STATE_STRIDE * queryIdx + componentIdx] = src;

			// compute length 
			if (componentIdx == 3)
			{
				float theta = distanceSO3[threadIdx.y];
				so3d[threadIdx.y] = 1.0 / sinf(theta);
				so3s0[threadIdx.y] = sinf((1.0 - ratio[threadIdx.y]) * theta);
				so3s1[threadIdx.y] = sinf(ratio[threadIdx.y] * theta);
			}

			so3Inter[sharedDataIdx] = src * dest;
			__syncthreads();

			if (componentIdx == 3)
			{
				float dq = so3Inter[sharedDataIdx] + so3Inter[sharedDataIdx + 1] + so3Inter[sharedDataIdx + 2] + so3Inter[sharedDataIdx + 3];
				if (dq < 0)
					so3s1[threadIdx.y] = -so3s1[threadIdx.y];
			}
			__syncthreads();

			intersectionStates[STATE_STRIDE * queryIdx + componentIdx] = (src * so3s0[threadIdx.y] + dest * so3s1[threadIdx.y]) * so3d[threadIdx.y];
		}

		if (componentIdx == 0)
		{
			// TODO: debug
			__syncthreads();

			printf("%d intersection state from (%.2f %.2f %.2f) to (%.2f %.2f %.2f) : \n(%.2f %.2f %.2f)\n", queryIdx,
				srcState[STATE_STRIDE * threadIdx.y + 0],
				srcState[STATE_STRIDE * threadIdx.y + 1],
				srcState[STATE_STRIDE * threadIdx.y + 2],
				destState[STATE_STRIDE * threadIdx.y + 0],
				destState[STATE_STRIDE * threadIdx.y + 1],
				destState[STATE_STRIDE * threadIdx.y + 2],
				intersectionStates[STATE_STRIDE * queryIdx + 0],
				intersectionStates[STATE_STRIDE * queryIdx + 1],
				intersectionStates[STATE_STRIDE * queryIdx + 2]);
		}
	}
}

void computeIntersectionStates(int queryWidth, float* treeNodes, int* nodeIndices, float* destStates, float* distR3, float* distSO3,
	float* intersectionStates, int* treeNodeSampleIndices)
{
	int gridX = (queryWidth + 31) / 32;
	dim3 grids = dim3(gridX, 1, 1);
	dim3 threads = dim3(7, 32, 1);

	computeIntersectionStatesKernel <<< grids, threads >>> (queryWidth, treeNodes, nodeIndices, destStates, distR3, distSO3, 
		intersectionStates,	treeNodeSampleIndices, d_levels);
}

__global__ void initRandomStates(int numStates, curandState* state)
{
	int id = threadIdx.y * blockDim.x + threadIdx.x;

	curandState* localState = state + id;
	curand_init(RANDOM_SEED, id, 0, localState);
}

__global__ void generateRandomStates(int numStates, curandState* state, float* randomStates)
{
	__shared__ float rotationalRandomValues[M * 3];

	int id = threadIdx.y * blockDim.x + threadIdx.x;
	int stateIdx = threadIdx.x;
	int componentIdx = threadIdx.y;

	curandState* localState = state + id;

	// translational value
	float x = curand_uniform(localState);
	randomStates[stateIdx * STATE_STRIDE + componentIdx] = x * c_boundsDiff[componentIdx] + c_boundsMin[componentIdx];

	// rotational value (quaternion)
	x = curand_uniform(localState);
	rotationalRandomValues[id] = x;
	__syncthreads();

	if (componentIdx < 2)
	{
		float x1 = rotationalRandomValues[stateIdx * 3];
		float c1 = sqrtf(1 - x1 + (-1 + 2 * x1) * componentIdx);
		float x2 = 2 * rotationalRandomValues[stateIdx * 3 + componentIdx + 1];
		float sinpix2, cospix2;
		sincospif(x2, &sinpix2, &cospix2);
		randomStates[stateIdx * STATE_STRIDE + 3 + componentIdx * 2] = sinpix2 * c1;
		randomStates[stateIdx * STATE_STRIDE + 4 + componentIdx * 2] = cospix2 * c1;
	}

	// debug
	if (componentIdx == 0)
	{
		printf("%d : %.3f %.3f %.3f %.3f %.3f %.3f %.3f\n", stateIdx, randomStates[stateIdx * STATE_STRIDE + 0],
			randomStates[stateIdx * STATE_STRIDE + 1],
			randomStates[stateIdx * STATE_STRIDE + 2],
			randomStates[stateIdx * STATE_STRIDE + 3],
			randomStates[stateIdx * STATE_STRIDE + 4],
			randomStates[stateIdx * STATE_STRIDE + 5],
			randomStates[stateIdx * STATE_STRIDE + 6]);
	}
}

__global__ void computeDistanceToDestKernel(int queryWidth, int k, float* samples, int* inIndices, float* inDistances, float* destStates, 
	float* outDistances, int* levels)
{
	__shared__ float sharedNeighbors[STATE_STRIDE * 32];
	__shared__ float sharedDest[STATE_STRIDE * 32];
	__shared__ float sharedDistancesR3[32];
	__shared__ float sharedDistancesSO3[32];

	int queryId = blockIdx.x * 32 + threadIdx.y;
	int kId = blockIdx.y;

	if (queryId < queryWidth)
	{
		copyIndexedStatesFromGlobalToShared(queryWidth, samples, inIndices + kId * queryWidth, sharedNeighbors);

		int sharedDataIdx = threadIdx.y * STATE_STRIDE + threadIdx.x;
		sharedDest[sharedDataIdx] = destStates[STATE_STRIDE * queryId + threadIdx.x];

		__syncthreads();

		computeDistances(queryWidth, sharedNeighbors, sharedDest, sharedDistancesR3, sharedDistancesSO3);

		if (threadIdx.x == 0)
		{
			float distToIntersect = inDistances[kId * queryWidth + queryId];
			float dist = sharedDistancesR3[threadIdx.y] + sharedDistancesSO3[threadIdx.y];
			if (distToIntersect > c_radius[levels[inIndices[kId * queryWidth + queryId]]])
				dist = 1E+37;
			outDistances[queryId * k + kId] = dist;

			printf("%d(%d) %d dist : %f\n", 
				queryId,
				(inIndices + kId * queryWidth)[queryId],
				kId,
				outDistances[queryId * k + kId]);
		}
	}
}

void computeDistanceToDest(int queryWidth, int k, float* samples, int* inIndices, float* inDistances, float* destStates, float* outDistances)
{
	int gridX = (queryWidth + 31) / 32;
	dim3 grids = dim3(gridX, k, 1);
	dim3 threads = dim3(7, 32, 1);

	computeDistanceToDestKernel <<< grids, threads >>> (queryWidth, k, samples, inIndices, inDistances, destStates, outDistances, d_levels);
}

__global__ void computeIndexedIndices(int queryWidth, int* inIndices, int* indicesOfIndices, int* outIndices)
{
	int queryIdx = blockIdx.x * 512 + threadIdx.x;
	outIndices[queryIdx] = inIndices[queryIdx + indicesOfIndices[queryIdx] * queryWidth];

	printf("Index chosen for %d : %d\n", queryIdx, outIndices[queryIdx]);
}

void computeMinDistanceSamples(int queryWidth, int k, float* samples, int *inIndices, float* inDistances, int* outIndices, float* outDistances)
{
	int blockSize = 512;
	while (blockSize / 2 >= k)
	{                                                                        
		blockSize /= 2;
	}
	int gridX = (k + (blockSize-1)) / blockSize;
	dim3 grids = dim3(gridX, queryWidth, 1);
	dim3 threads = dim3(blockSize, 1, 1);
	int smemSize = blockSize * (sizeof(float) + sizeof(int));

	switch (blockSize)
	{
	case 512:
		reduceMin<256, float><<< grids, threads, smemSize >>>(inDistances, outDistances, k, queryWidth, outIndices, false); break;
	case 256:
		reduceMin<128, float><<< grids, threads, smemSize >>>(inDistances, outDistances, k, queryWidth, outIndices, false); break;
	case 128:
		reduceMin< 64, float><<< grids, threads, smemSize >>>(inDistances, outDistances, k, queryWidth, outIndices, false); break;
	case 64:
		reduceMin< 32, float><<< grids, threads, smemSize >>>(inDistances, outDistances, k, queryWidth, outIndices, false); break;
	case 32:
		reduceMin< 16, float><<< grids, threads, smemSize >>>(inDistances, outDistances, k, queryWidth, outIndices, false); break;
	case 16:
		reduceMin< 8, float><<< grids, threads, smemSize >>>(inDistances, outDistances, k, queryWidth, outIndices, false); break;
	case 8:
		reduceMin< 4, float><<< grids, threads, smemSize >>>(inDistances, outDistances, k, queryWidth, outIndices, false); break;
	case 4:
		reduceMin< 2, float><<< grids, threads, smemSize >>>(inDistances, outDistances, k, queryWidth, outIndices, false); break;
	case 2:
		reduceMin< 1, float><<< grids, threads, smemSize >>>(inDistances, outDistances, k, queryWidth, outIndices, false); break;
	case 1:
		reduceMin1<<< grids, threads, smemSize >>>(inDistances, outDistances, queryWidth, outIndices, false); break;
	}

	{
		int gridX = (queryWidth + (512-1)) / 512;
		dim3 grids = dim3(gridX, 1, 1);
		dim3 threads = dim3(min(queryWidth, 512), 1, 1);

		computeIndexedIndices <<< grids, threads >>>(queryWidth, inIndices, outIndices, outIndices);
	}
	
}

template <class CollideFunctor>
__global__ void checkCollisionKernel(int queryWidth, float* treeNodes, int* inTreeIndices, float* samples, int* inSampleIndices, float* distances, 
	GMesh* robot, GMesh* obstacle, CollideFunctor collideF,
	float stateValidityCheckingResolutionR3, float stateValidityCheckingResolutionSO3, int* outResult, float* colFreeStates)
{
	__shared__ float sharedFromStates[32 * STATE_STRIDE];
	__shared__ float sharedToStates[32 * STATE_STRIDE];
	__shared__ float sharedDistancesR3[32];
	__shared__ float sharedDistancesSO3[32];
	__shared__ float sharedN[32];
	__shared__ float sharedInterpolatedStates[32 * STATE_STRIDE];

	int queryId = blockIdx.x * 32 + threadIdx.y;
	int globalDataIdx = queryId * STATE_STRIDE + threadIdx.x;
	int sharedDataIdx = threadIdx.y * STATE_STRIDE + threadIdx.x;

	if (queryId < queryWidth)
	{
		copyIndexedStatesFromGlobalToShared(queryWidth, treeNodes, inTreeIndices, sharedFromStates);
		copyIndexedStatesFromGlobalToShared(queryWidth, samples, inSampleIndices, sharedToStates);

		computeDistances(queryWidth, sharedFromStates, sharedToStates, sharedDistancesR3, sharedDistancesSO3);
		float distance = sharedDistancesR3[threadIdx.y] + sharedDistancesSO3[threadIdx.y];

		if (threadIdx.x == 0)
		{
			sharedN[threadIdx.y] = (int)ceilf(max(sharedDistancesR3[threadIdx.y] * stateValidityCheckingResolutionR3, 
				sharedDistancesSO3[threadIdx.y] * stateValidityCheckingResolutionSO3));
		}
		__syncthreads();

		float prevInterp = sharedFromStates[sharedDataIdx];
		int result = COLLISION_FREE;
		for (int i = 1; i <= sharedN[threadIdx.y]; ++i)
		{
			float t = (float)i / sharedN[threadIdx.y];
			interpolateState(queryWidth, sharedFromStates, sharedToStates, t, sharedDistancesSO3, sharedInterpolatedStates);
			__syncthreads();

			if (threadIdx.x == 0)
			{
				printf("Interp from (%.3f %.3f %.3f) to (%.3f %.3f %.3f) at %.3f:\n(%.3f %.3f %.3f)\n",
					sharedFromStates[STATE_STRIDE * threadIdx.y + 0],
					sharedFromStates[STATE_STRIDE * threadIdx.y + 1],
					sharedFromStates[STATE_STRIDE * threadIdx.y + 2],
					sharedToStates[STATE_STRIDE * threadIdx.y + 0],
					sharedToStates[STATE_STRIDE * threadIdx.y + 1],
					sharedToStates[STATE_STRIDE * threadIdx.y + 2],
					t,
					sharedInterpolatedStates[STATE_STRIDE * threadIdx.y + 0],
					sharedInterpolatedStates[STATE_STRIDE * threadIdx.y + 1],
					sharedInterpolatedStates[STATE_STRIDE * threadIdx.y + 2]);
			}
			__syncthreads();

			float* state = sharedInterpolatedStates + STATE_STRIDE * threadIdx.y;
			result = collideF(robot->models, robot->vertexPointers, robot->triIdxPointers, state,
				obstacle->models, obstacle->vertexPointers, obstacle->triIdxPointers);

			if (result == IN_COLLISION)
			{
				if (i == 1)
					result = IN_COLLISION_WITH_NO_ADVANCE;

				colFreeStates[globalDataIdx] = prevInterp;
				break;
			}
			prevInterp = sharedInterpolatedStates[sharedDataIdx];
		}
		
		colFreeStates[globalDataIdx] = prevInterp;
		if (threadIdx.x == 0)
		{
			outResult[queryId] = result;

			printf("Collision Checking from (%.3f %.3f %.3f | %.3f %.3f %.3f %.3f) to (%.3f %.3f %.3f | %.3f %.3f %.3f %.3f) : %d\n",
				sharedFromStates[threadIdx.y * STATE_STRIDE + 0],
				sharedFromStates[threadIdx.y * STATE_STRIDE + 1],
				sharedFromStates[threadIdx.y * STATE_STRIDE + 2],
				sharedFromStates[threadIdx.y * STATE_STRIDE + 3],
				sharedFromStates[threadIdx.y * STATE_STRIDE + 4],
				sharedFromStates[threadIdx.y * STATE_STRIDE + 5],
				sharedFromStates[threadIdx.y * STATE_STRIDE + 6],
				sharedToStates[threadIdx.y * STATE_STRIDE + 0],
				sharedToStates[threadIdx.y * STATE_STRIDE + 1],
				sharedToStates[threadIdx.y * STATE_STRIDE + 2],
				sharedToStates[threadIdx.y * STATE_STRIDE + 3],
				sharedToStates[threadIdx.y * STATE_STRIDE + 4],
				sharedToStates[threadIdx.y * STATE_STRIDE + 5],
				sharedToStates[threadIdx.y * STATE_STRIDE + 6],
				result);
			__syncthreads();
		}
	}
}

void checkCollision(int queryWidth, float* treeNodes, int* inTreeIndices, float* samples, int* inSampleIndices, int* outResult, float* colFreeStates,
	float stateValidityCheckingResolutionR3, float stateValidityCheckingResolutionSO3)
{
	int* d_N = d_internalValues;
	int* d_MaxN = d_internalValues2;

	// compute from/to states
	int gridX = (queryWidth + 31) / 32;
	dim3 grids = dim3(gridX, 1, 1);
	dim3 threads = dim3(7, 32, 1);

	// compute N (# of checking points)
	BVHStateCollideFunctor<OBBNode, OBB, ushort2, 1> collideF;
	checkCollisionKernel<BVHStateCollideFunctor<OBBNode, OBB, ushort2, 1> > <<< grids, threads >>>(
		queryWidth, treeNodes, inTreeIndices, samples, inSampleIndices, d_distances, 
		d_robot, d_obstacle, collideF,
		stateValidityCheckingResolutionR3, stateValidityCheckingResolutionSO3, outResult, colFreeStates);
}

__global__ void addNewStatesKernel(int queryWidth, float* samples, int* fromTreeIndices, int* toSampleIndices,
	float* treeNodes, int* nodeIndices, int* parentIndices, unsigned int* numNodes, float* colFreeStates, int* colFreeResult)
{
	__shared__ int sharedCollisionResult[32];
	__shared__ int sharedParentTreeIndex[32];
	__shared__ float sharedColFreeStates[32 * STATE_STRIDE];
	__shared__ unsigned int sharedNewNodeId[32];
	__shared__ unsigned int sharedNewSampleId[32];

	int queryId = blockIdx.x * 32 + threadIdx.y;
	if (queryId < queryWidth)
	{
		if (threadIdx.x == 0)
		{
			sharedCollisionResult[threadIdx.y] = colFreeResult[queryId];
			sharedParentTreeIndex[threadIdx.y] = fromTreeIndices[queryId];

			if (sharedCollisionResult[threadIdx.y] != IN_COLLISION_WITH_NO_ADVANCE)
				sharedNewNodeId[threadIdx.y] = atomicInc(numNodes, UINT_MAX);
			if (sharedCollisionResult[threadIdx.y] == IN_COLLISION)
				sharedNewSampleId[threadIdx.y] = atomicInc(&d_numSamples, UINT_MAX);
		}
		sharedColFreeStates[threadIdx.y * STATE_STRIDE + threadIdx.x] = colFreeStates[threadIdx.y * STATE_STRIDE + threadIdx.x];
		__syncthreads();

		if (sharedCollisionResult[threadIdx.y] != IN_COLLISION_WITH_NO_ADVANCE)
		{
			if (threadIdx.x == 0)
			{
				parentIndices[sharedNewNodeId[threadIdx.y]] = fromTreeIndices[queryId];
				if (sharedCollisionResult[threadIdx.y] == IN_COLLISION)
				{
					nodeIndices[sharedNewNodeId[threadIdx.y]] = sharedNewSampleId[threadIdx.y];
					// TODO:
					//d_isTreeNode[sharedNewSampleId[threadIdx.y]] = 1;
					//d_levels[sharedNewSampleId[threadIdx.y]] = d_levels[nodeIndices[fromTreeIndices[queryId]]];
				}
				else // collision free
					nodeIndices[sharedNewNodeId[threadIdx.y]] = toSampleIndices[queryId];
				
			}
			treeNodes[STATE_STRIDE * sharedNewNodeId[threadIdx.y] + threadIdx.x] = sharedColFreeStates[threadIdx.y * STATE_STRIDE + threadIdx.x];
			if (sharedCollisionResult[threadIdx.y] == IN_COLLISION)
				samples[STATE_STRIDE * sharedNewSampleId[threadIdx.y] + threadIdx.x] = sharedColFreeStates[threadIdx.y * STATE_STRIDE + threadIdx.x];
		}
	}
}

void addNewStates(int queryWidth, float* samples, int* fromTreeIndices, int* toSampleIndices,
	float* treeNodes, int* nodeIndices, int* parentIndices, unsigned int* numNodes, float* colFreeStates, int* colFreeResult)
{
	int gridX = (queryWidth + 31) / 32;
	dim3 grids = dim3(gridX, 1, 1);
	dim3 threads = dim3(7, 32, 1);

	addNewStatesKernel<<<grids, threads>>>(queryWidth, samples, fromTreeIndices, toSampleIndices, 
		treeNodes, nodeIndices, parentIndices, numNodes, colFreeStates, colFreeResult);
}

namespace poisson
{
	GPUPoissonCuda::GPUPoissonCuda()
		: resetRequired_(false), stateValidityCheckingResolutionR3_(0.01f), stateValidityCheckingResolutionSO3_(0.01f)
	{

	}

	GPUPoissonCuda::~GPUPoissonCuda()
	{

	}

	void GPUPoissonCuda::setup(const std::vector<float>& vBoundsMin, const std::vector<float>& vBoundsMax, 
		const ompl::app::GeometrySpecification* geometrySpecification,
		float stateValidityCheckingResolutionR3, float stateValidityCheckingResolutionSO3)
	{
		initializeMemory();
		uploadData(vBoundsMin, vBoundsMax);

		stateValidityCheckingResolutionR3_ = stateValidityCheckingResolutionR3;
		stateValidityCheckingResolutionSO3_ = stateValidityCheckingResolutionSO3;

		GMesh* h_obstacles;
		GMesh* h_robot;
		gplanner_collision_wrapper::setup(geometrySpecification, &h_obstacles, &h_robot);
		TOGPU(d_robot, h_robot, sizeof(GMesh));
		TOGPU(d_obstacle, h_obstacles, sizeof(GMesh));



		dim3 grids = dim3(1, 1, 1);
		dim3 threads = dim3(M, 3, 1);
		initRandomStates <<< grids, threads >>> (M, d_randStates);
	}

	void GPUPoissonCuda::solve()
	{
		resetRequired_ = true;

		bool terminate = false;
		bool expandStartTree = true;
		while (!terminate)
		{
			// generate M random configurations
			generateMRandomStates();

			// compute nearest nodes
			computeNearestNodes(expandStartTree);

			// compute neighbor sample
			computeNeighborSample(expandStartTree);


			// generate new sample if no neighbor exists
			// TODO?

			// collision check
			checkCollision(expandStartTree);

			// if no collision

			//   connect

			// if collision 

			//   connect to last collision-free state

			//   shrink sample


			//   compute nearest nodes in other tree

			//   compute neighbor sample

			//   generate new sample if no neighbor exists

			//   collision check

			//   if no collision

			//     if is current tree node -> solution found

			//     else connect

			//   if collision

			//   connect to last collision-free state

			//   shrink sample

		}
	}

	void GPUPoissonCuda::reset()
	{
		// TODO: reset?

		resetRequired_ = false;
	}

	void GPUPoissonCuda::addStartGoalPosition(const float* startGoalPos)
	{
		TOGPU(d_samples + h_numPrecomputedSamples * STATE_STRIDE, startGoalPos, STATE_SIZE_STRIDE * 2);
		GPUTOGPU(d_startTreeNodes, d_samples + h_numPrecomputedSamples * STATE_STRIDE, STATE_SIZE_STRIDE);
		GPUTOGPU(d_goalTreeNodes, d_samples + (h_numPrecomputedSamples + 1) * STATE_STRIDE, STATE_SIZE_STRIDE);

		const int isTreeNode[2] = {START_TREE, GOAL_TREE};
		TOGPU(d_isTreeNode + h_numPrecomputedSamples * STATE_STRIDE, isTreeNode, sizeof(int) * 2);

		h_numSamples = h_numPrecomputedSamples;
		TOGPU(d_startTreeNodeIndices, &h_numSamples, sizeof(int));
		++h_numSamples;
		TOGPU(d_goalTreeNodeIndices, &h_numSamples, sizeof(int));
		++h_numSamples;

		const int noParent = -1;
		TOGPU(d_startTreeNodeParentIndices, &noParent, sizeof(int));
		TOGPU(d_goalTreeNodeParentIndices, &noParent, sizeof(int));

		h_numStartTreeNodes = 1;
		h_numGoalTreeNodes = 1;

		TOGPU_CONSTANT(d_numSamples, &h_numSamples, sizeof(unsigned int));
		TOGPU(&d_numStartTreeNodes, &h_numStartTreeNodes, sizeof(unsigned int));
		TOGPU(&d_numGoalTreeNodes, &h_numGoalTreeNodes, sizeof(unsigned int));
	}

	void GPUPoissonCuda::initializeMemory()
	{
		GPUMALLOC(&d_samples, STATE_SIZE_STRIDE * MAX_SAMPLES);
		GPUMALLOC(&d_levels, sizeof(int) * MAX_SAMPLES);
		GPUMALLOC(&d_isTreeNode, sizeof(int) * MAX_SAMPLES);

		GPUMALLOC(&d_numStartTreeNodes, sizeof(unsigned int));
		GPUMALLOC(&d_startTreeNodeIndices, sizeof(int) * MAX_NODES);
		GPUMALLOC(&d_startTreeNodeParentIndices, sizeof(int) * MAX_NODES);
		GPUMALLOC(&d_startTreeNodes, STATE_SIZE_STRIDE * MAX_NODES);
		d_startPos = d_startTreeNodes;

		GPUMALLOC(&d_numGoalTreeNodes, sizeof(unsigned int));
		GPUMALLOC(&d_goalTreeNodeIndices, sizeof(int) * MAX_NODES);
		GPUMALLOC(&d_goalTreeNodeParentIndices, sizeof(int) * MAX_NODES);
		GPUMALLOC(&d_goalTreeNodes, STATE_SIZE_STRIDE * MAX_NODES);
		d_goalPos = d_goalTreeNodes + STATE_STRIDE;

		GPUMALLOC(&d_internalStates, STATE_SIZE_STRIDE * M);
		GPUMALLOC(&d_internalStates2, STATE_SIZE_STRIDE * M);

		GPUMALLOC(&d_randStates, sizeof(curandState) * M * 3);

		GPUMALLOC(&d_distances, sizeof(float) * (MAX_K + 512) * M);
		GPUMALLOC(&d_nnIndices, sizeof(int) * MAX_K * M);
		GPUMALLOC(&d_internalDistances, sizeof(float) * M * MAX_SAMPLES);

		GPUMALLOC(&d_nnIndices2, sizeof(int) * MAX_K * M);
		GPUMALLOC(&d_nnIndices3, sizeof(int) * (MAX_K + 512) * M);

		GPUMALLOC(&d_result, sizeof(int) * M);
		
		GPUMALLOC(&d_internalValues, sizeof(float) * M);
		GPUMALLOC(&d_internalValues2, sizeof(float) * M);

		GPUMALLOC(&d_robot, sizeof(GMesh));
		GPUMALLOC(&d_obstacle, sizeof(GMesh));
	}

	void GPUPoissonCuda::freeMemory()
	{
		GPUFREE(d_samples);
		
		GPUFREE(d_levels);
		GPUFREE(d_isTreeNode);

		GPUFREE(d_numStartTreeNodes);
		GPUFREE(d_startTreeNodeIndices);
		GPUFREE(d_startTreeNodeParentIndices);
		GPUFREE(d_startTreeNodes);

		GPUFREE(d_numGoalTreeNodes);
		GPUFREE(d_goalTreeNodeIndices);
		GPUFREE(d_goalTreeNodeParentIndices);
		GPUFREE(d_goalTreeNodes);

		GPUFREE(d_internalStates);
		GPUFREE(d_internalStates2);

		GPUFREE(d_randStates);

		GPUFREE(d_distances);
		GPUFREE(d_nnIndices);
		GPUFREE(d_internalDistances);

		GPUFREE(d_nnIndices2);
		GPUFREE(d_nnIndices3);

		GPUFREE(d_result);

		GPUFREE(d_internalValues);
		GPUFREE(d_internalValues2);

		GPUFREE(d_robot);
		GPUFREE(d_obstacle);

		CUDA_SAFE_CALL(cudaDeviceReset());
	}

	void GPUPoissonCuda::uploadData(const std::vector<float>& vBoundsMin, const std::vector<float>& vBoundsMax)
	{
		h_numPrecomputedSamples = states_.size();

		float* h_samples;
		CPUMALLOC(&h_samples, STATE_SIZE_STRIDE * h_numPrecomputedSamples);

		for (unsigned int i = 0; i < h_numPrecomputedSamples; ++i)
		{
			const double* stateValues = states_[i].getValue();

			for (int j = 0; j < 7; ++j)
			{
				*(h_samples + STATE_STRIDE * i + j) = (float)(stateValues[j]);
			}
		}
		TOGPU(d_samples, h_samples, STATE_SIZE_STRIDE * h_numPrecomputedSamples);

		CPUFREE(h_samples);

		GPUMEMSET(d_levels, 0, sizeof(int) * MAX_SAMPLES);
		GPUMEMSET(d_isTreeNode, 0, sizeof(int) * MAX_SAMPLES);

		TOGPU_CONSTANT(c_boundsMin, &(vBoundsMin[0]), sizeof(float) * 3);
		TOGPU_CONSTANT(c_boundsMax, &(vBoundsMax[0]), sizeof(float) * 3);

		std::vector<float> vBoundsDiff(3);
		for (int i = 0; i < 3; ++i)
			vBoundsDiff[i] = vBoundsMax[i] - vBoundsMin[i];
		TOGPU_CONSTANT(c_boundsDiff, &(vBoundsDiff[0]), sizeof(float) * 3);

		std::vector<float> vRadius(MAX_LEVEL);
		vRadius[0] = BASE_RADIUS;
		for (int i = 1; i < MAX_LEVEL; ++i)
			vRadius[i] = vRadius[i - 1] * 0.5f;
		TOGPU_CONSTANT(c_radius, &(vRadius[0]), sizeof(float) * MAX_LEVEL);
	}

	void GPUPoissonCuda::generateMRandomStates()
	{
		dim3 grids = dim3(1, 1, 1);
		dim3 threads = dim3(M, 3, 1);

		generateRandomStates <<< grids, threads >>> (M, d_randStates, d_internalStates);

		// out : d_internalStates has generated random states
	}

	void GPUPoissonCuda::computeNearestNodes(bool startTree)
	{
		// source : randomStates
		// ref : startTreeNodes / goalTreeNodes
		knn(d_startTreeNodes, h_numStartTreeNodes, d_internalStates, M, 1, d_distances, d_nnIndices);

		// out: d_nnIndices has tree indices of selected tree nodes
	}

	void GPUPoissonCuda::computeNeighborSample(bool expandStartTree)
	{
		// get distances from tree nodes to random states
		computeDistancesTreeNodesToStates(M, d_startTreeNodes, d_nnIndices, d_internalStates, d_distances, d_distances + M);

		// compute intersection states
		computeIntersectionStates(M, d_startTreeNodes, d_nnIndices, d_internalStates, d_distances, d_distances + M,
			d_internalStates2, d_startTreeNodeIndices);

		// get knn poisson samples near intersection states in the disk radius
		unsigned int k = __min(MAX_K, h_numSamples);
		knn(d_samples, h_numSamples, d_internalStates2, M, k, d_distances, d_nnIndices3);

		float* h_samples = new float[STATE_STRIDE*h_numSamples];
		float* h_internalStates2 = new float[STATE_STRIDE*M];
		float* h_inDistances = new float[k*M];
		int* h_inIndices = new int[k*M];

		FROMGPU(h_samples, d_samples, sizeof(float)*STATE_STRIDE*h_numSamples);
		FROMGPU(h_internalStates2, d_internalStates2, sizeof(float)*M);
		FROMGPU(h_inDistances, d_distances, sizeof(float)*k*M);
		FROMGPU(h_inIndices, d_nnIndices3, sizeof(int)*k*M);
		cudaThreadSynchronize();

		delete[] h_samples;
		delete[] h_internalStates2;
		delete[] h_inDistances;
		delete[] h_inIndices;

		// compute distances to random states
		computeDistanceToDest(M, k, d_samples, d_nnIndices3, d_distances, d_internalStates, d_internalDistances);

		// get the closest(nearest neighbor) samples to random states
		computeMinDistanceSamples(M, k, d_samples, d_nnIndices3, d_internalDistances, d_nnIndices2, d_distances);

		// out: d_nnIndices2 has state indices of selected poisson states
	}

	void GPUPoissonCuda::checkCollision(bool expandStartTree)
	{
		float* d_colFreeStates = d_internalStates2;
		int* d_colFreeResult = d_result;
		::checkCollision(M, d_startTreeNodes, d_nnIndices, d_samples, d_nnIndices2, d_result, d_colFreeStates, 
			stateValidityCheckingResolutionR3_, stateValidityCheckingResolutionSO3_);

		// out: d_result has collision result (0: collision-free, 1: collision detected)
		//      d_internalStates2 has last collision-free states

		// add collision-free states
		addNewStates(M, d_samples, d_nnIndices, d_nnIndices2, 
			d_startTreeNodes, d_startTreeNodeIndices, d_startTreeNodeParentIndices, d_numStartTreeNodes, d_colFreeStates, d_colFreeResult);

		// shrink node if collision is detected
	}
}
